{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#hide\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp core"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "from nbdev.export import check_re,read_nb\n",
    "from pathlib import Path\n",
    "import re\n",
    "import os\n",
    "import platform"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# core\n",
    "\n",
    "> contains all the functions of fast_neptune"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Code cells"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here we define regex, inspired by the ones used in `nbdev`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "_re_blank_code = re.compile(r\"\"\"\n",
    "# Matches any line with #export or #exports without any module name:\n",
    "^         # beginning of line (since re.MULTILINE is passed)\n",
    "\\s*       # any number of whitespace\n",
    "\\#\\s*     # # then any number of whitespace\n",
    "code  # export or exports\n",
    "\\s*       # any number of whitespace\n",
    "$         # end of line (since re.MULTILINE is passed)\n",
    "\"\"\", re.IGNORECASE | re.MULTILINE | re.VERBOSE)\n",
    "\n",
    "_re_mod_code = re.compile(r\"\"\"\n",
    "# Matches any line with #export or #exports with a module name and catches it in group 1:\n",
    "^         # beginning of line (since re.MULTILINE is passed)\n",
    "\\s*       # any number of whitespace\n",
    "\\#\\s*     # # then any number of whitespace\n",
    "code  # export or exports\n",
    "\\s*       # any number of whitespace\n",
    "(\\S+)     # catch a group with any non-whitespace chars\n",
    "\\s*       # any number of whitespace\n",
    "$         # end of line (since re.MULTILINE is passed)\n",
    "\"\"\", re.IGNORECASE | re.MULTILINE | re.VERBOSE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def is_code(cell: dict, default=\"main.py\") -> str:\n",
    "    \"\"\"Checks if `cell` is to be exported and returns the name of the module to export it if provided.\n",
    "    \n",
    "    Parses the cell of a Jupyter Notebook, checks if it has the #code tag,\n",
    "    and returns the name of the module it is associated with, otherwise returns\n",
    "    None.\n",
    "    \n",
    "    Args:\n",
    "        cell: dict of the JSON of the cell\n",
    "        default: name of the default module where all #code will be added\n",
    "    Returns:\n",
    "        The string of the name of the cell it is associated if #code is present,\n",
    "        else None\n",
    "    \"\"\"\n",
    "    if check_re(cell, _re_blank_code):\n",
    "        return default\n",
    "    tst = check_re(cell, _re_mod_code)\n",
    "    return os.path.sep.join(tst.groups()[0].split('.')).replace(\"\\\\\",\".\").replace(\"/\",\".\") if tst else None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "from collections import defaultdict\n",
    "\n",
    "def get_codes(fn:str,default:str = \"main.py\") -> dict:\n",
    "    \"\"\"Returns a dictionary where each key contains the name\n",
    "    of the module, and the value is the code in it.\"\"\"\n",
    "    nb = read_nb(fn)\n",
    "    \n",
    "    module_to_code = defaultdict(str)\n",
    "    \n",
    "    module_to_code[default] = \"\"\n",
    "    \n",
    "    for cell in nb[\"cells\"]:\n",
    "        code = is_code(cell,default)\n",
    "        if code:\n",
    "            module_to_code[code] += cell[\"source\"]\n",
    "    \n",
    "    return dict(module_to_code)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Properties"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### OS information"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def get_metadata() -> dict:\n",
    "    \"\"\"Returns metadata about the current running environment.\"\"\"\n",
    "    data = {\n",
    "        \"os\":os.name,\n",
    "        \"system\":platform.system(),\n",
    "        \"release\":platform.release(),\n",
    "        \"python_version\":platform.python_version()\n",
    "    }\n",
    "    return data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'os': 'nt', 'system': 'Windows', 'release': '10', 'python_version': '3.7.5'}"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "get_metadata()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Modules information"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def create_requirements(fn):\n",
    "    \"\"\"Create requirements file\"\"\"\n",
    "    # Convert the notebook to a python file\n",
    "    os.system(f\"jupyter nbconvert --to=python {fn}\")\n",
    "    \n",
    "    # Create the requirements file\n",
    "    os.system(\"pipreqs ./ --force\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Properties cells"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "_re_blank_property = re.compile(r\"\"\"\n",
    "# Matches any line with #export or #exports without any module name:\n",
    "^         # beginning of line (since re.MULTILINE is passed)\n",
    "\\s*       # any number of whitespace\n",
    "\\#\\s*     # # then any number of whitespace\n",
    "property  # export or exports\n",
    "\\s*       # any number of whitespace\n",
    "$         # end of line (since re.MULTILINE is passed)\n",
    "\"\"\", re.IGNORECASE | re.MULTILINE | re.VERBOSE)\n",
    "\n",
    "_re_obj_def = re.compile(r\"\"\"\n",
    "# Catches any 0-indented object definition (bla = thing) with its name in group 1\n",
    "^          # Beginning of a line (since re.MULTILINE is passed)\n",
    "([^=\\s]*)  # Catching group with any character except a whitespace or an equal sign\n",
    "\\s*=       # Any number of whitespace followed by an =\n",
    "\"\"\", re.MULTILINE | re.VERBOSE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def is_property(cell):\n",
    "    \"Check if `cell` is to be exported and returns the name of the module to export it if provided\"\n",
    "    if check_re(cell, _re_blank_property):\n",
    "        return True\n",
    "    else:\n",
    "        return False\n",
    "\n",
    "def add_cell_to_properties(cell: dict,properties: dict,globs:dict):\n",
    "    \"\"\"Adds all variables in the cell to the properties\"\"\"\n",
    "    objs = _re_obj_def.findall(cell[\"source\"])\n",
    "    \n",
    "    objs = {obj : globs[obj] for obj in objs}\n",
    "    \n",
    "    properties.update(objs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def files_in_properties(properties:dict):\n",
    "    \"\"\"Returns the list of files from properties\"\"\"\n",
    "    files = []\n",
    "    for key,val in properties.items():\n",
    "        if isinstance(val,Path) and val.is_file():\n",
    "            files.append(str(val))\n",
    "    return files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def get_properties_from_cells(fn: str,globs:dict,return_files:bool = True,):\n",
    "    \"\"\"Gets the properties from all #property cells\"\"\"\n",
    "    \n",
    "    nb = read_nb(fn)\n",
    "    \n",
    "    properties = {}\n",
    "    \n",
    "    for cell in nb[\"cells\"]:\n",
    "        if is_property(cell):\n",
    "            add_cell_to_properties(cell,properties,globs=globs)\n",
    "\n",
    "    files = files_in_properties(properties)\n",
    "    return properties,files"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Wrapper"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "from contextlib import contextmanager\n",
    "from neptune.projects import Project\n",
    "from neptune.experiments import Experiment\n",
    "\n",
    "@contextmanager\n",
    "def fast_experiment(project: Project,nb_name:str,globs:dict,return_files: bool = True,\n",
    "                    default:str = \"main.py\",**kwargs) -> Experiment:\n",
    "    \"\"\"Creates a Neptune ML experiment, wrapped with meta data.\n",
    "    \n",
    "    Args:\n",
    "        project: Neptune Project\n",
    "        nb_name: str name of the current notebook to be recorded\n",
    "        globs: dict of the global variables. Simply set globs = globals() and then pass it.\n",
    "        return_files: bool, True if we want to send files recorded in the parameters.\n",
    "        default: str name of the default code \n",
    "        kwargs: additional args passed to Neptune ML when the experiment is created\n",
    "        \n",
    "    Returns:\n",
    "        exp: Neptune ML experiment\n",
    "    \"\"\"\n",
    "    # First we get the code cells\n",
    "    codes = get_codes(nb_name,default=default)\n",
    "    \n",
    "    # We write them in separate files\n",
    "    for fn,code in codes.items():\n",
    "        with open(fn,\"w\") as file:\n",
    "            file.write(code)\n",
    "            \n",
    "    codes = list(codes.keys())\n",
    "    \n",
    "    # We get the properties\n",
    "    properties,files = get_properties_from_cells(nb_name,globs=globs,return_files=return_files)\n",
    "    metadata = get_metadata()\n",
    "    properties.update(metadata)\n",
    "    properties[\"nb_name\"] = nb_name\n",
    "    \n",
    "    # We convert the dict keys to string\n",
    "    for k,v in properties.items():\n",
    "        properties[k] = str(v)\n",
    "    \n",
    "    exp = project.create_experiment(params=properties,upload_source_files=codes,**kwargs)\n",
    "    \n",
    "    # We create the requirements file and send it\n",
    "    create_requirements(nb_name)\n",
    "    exp.send_artifact(\"requirements.txt\")\n",
    "    \n",
    "    for fn in files:\n",
    "        exp.send_artifact(fn)\n",
    "        \n",
    "    yield exp\n",
    "    \n",
    "    exp.stop()\n",
    "    \n",
    "    # We remove the code files\n",
    "    for fn in codes:\n",
    "        os.remove(fn)\n",
    "        \n",
    "    os.remove(\"requirements.txt\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
